from sentence_transformers import SentenceTransformer
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
import os
from flask import current_app
import numpy as np

class VectorSearch:
    def __init__(self, collection_name='order_addresses'):
        self.collection_name = collection_name
        self.model = SentenceTransformer(current_app.config['MODEL_NAME'])
        self.vector_dim = current_app.config['VECTOR_DIM']
        
        # 创建本地存储目录
        self.db_path = os.path.join(os.getcwd(), "milvus_db")
        os.makedirs(self.db_path, exist_ok=True)
        
        # 连接到Milvus
        self._connect_to_milvus()
        
        # 初始化集合
        self._init_collection()
    
    def _connect_to_milvus(self):
        """连接到Milvus服务"""
        try:
            connections.connect(
                alias="default",
                uri=f"lite://{self.db_path}"
            )
        except Exception as e:
            raise ConnectionError(f"连接Milvus失败: {str(e)}")

    def _init_collection(self):
        """初始化集合"""
        if utility.has_collection(self.collection_name):
            self.collection = Collection(self.collection_name)
            return
        
        fields = [
            FieldSchema(name="order_id", dtype=DataType.VARCHAR, max_length=36, is_primary=True),
            FieldSchema(name="user_id", dtype=DataType.VARCHAR, max_length=36),
            FieldSchema(name="address", dtype=DataType.VARCHAR, max_length=500),
            FieldSchema(name="address_vector", dtype=DataType.FLOAT_VECTOR, dim=self.vector_dim)
        ]
        
        schema = CollectionSchema(fields=fields, description="订单地址向量搜索")
        self.collection = Collection(name=self.collection_name, schema=schema)
        
        # 创建索引
        index_params = {
            "metric_type": "COSINE",
            "index_type": "IVF_FLAT",
            "params": {"nlist": 128}
        }
        self.collection.create_index(field_name="address_vector", index_params=index_params)

    def vectorize_addresses(self, orders):
        """向量化订单地址并存储"""
        try:
            # 准备数据
            order_ids = [order['order_id'] for order in orders]
            user_ids = [order['user_id'] for order in orders]
            addresses = [order['address'] for order in orders]
            
            # 向量化地址
            vectors = self.model.encode(addresses)
            
            # 插入数据
            entities = [
                order_ids,
                user_ids,
                addresses,
                vectors.tolist()
            ]
            
            self.collection.insert(entities)
            self.collection.flush()
            
            return len(orders), 0
        except Exception as e:
            return 0, len(orders)

    def search_similar_addresses(self, query_address, threshold=0.85, limit=10):
        """搜索相似地址"""
        try:
            # 向量化查询地址
            query_vector = self.model.encode([query_address])[0]
            
            # 加载集合
            self.collection.load()
            
            # 搜索
            search_params = {"metric_type": "COSINE", "params": {"nprobe": 10}}
            results = self.collection.search(
                data=[query_vector],
                anns_field="address_vector",
                param=search_params,
                limit=limit,
                output_fields=["order_id", "user_id", "address"]
            )
            
            # 处理结果
            similar_orders = []
            for hits in results:
                for hit in hits:
                    if hit.score >= threshold:
                        similar_orders.append({
                            'order_id': hit.entity.get('order_id'),
                            'user_id': hit.entity.get('user_id'),
                            'address': hit.entity.get('address'),
                            'similarity': round(hit.score, 4)
                        })
            
            return similar_orders
        except Exception as e:
            raise Exception(f"搜索失败: {str(e)}") 